#!/usr/bin/env python
import itk
import sys
import os
import imageTools.ITKUtils.io as io
from . import utils
import multiprocessing as mp

def cast2uc3(fname, start, end, fcasted):
    reader = itk.ImageFileReader.IUC3.New()
    reader.SetFileName(fname)
    img = reader.GetOutput()
    tf = itk.ThresholdImageFilter.IUC3.New()
    tf.SetInput(img)
    tf.SetLower(1)
    tf.SetUpper(10)
    ccf = itk.ConnectedComponentImageFilter.IUC3IUL3.New()
    ccf.SetInput(tf.GetOutput())
    rcf = itk.RelabelComponentImageFilter.IUL3IUS3.New()
    rcf.SetInput(ccf.GetOutput()) 
    tf2 = itk.ThresholdImageFilter.IUS3.New()
    tf2.SetInput(rcf.GetOutput())
    tf2.SetLower(start)
    tf2.SetUpper(end)
    io.writeImage(tf2.GetOutput(), fcasted)
def combine(fname1, fname2, fname= "dummy.nii.gz" ):
    if fname1 == 'None':
        return fname2
    img1 = io.readImage(fname1, imgMode = 'UShort' )
    img2 = io.readImage(fname2, imgMode = 'UShort')
    cf = itk.AddImageFilter.IUS3IUS3IUS3.New()
    cf.SetInput1(img1)
    cf.SetInput2(img2)
    io.writeImage(cf.GetOutput(), fname= fname)
    return fname

def otsu_seg(args):
    fname = args[0]
    fout = args[1]
    img = io.readImage(fname)
    of = itk.OtsuThresholdImageFilter.IF3IF3.New()
    of.SetInput(img)
    of.SetInsideValue(0)
    of.SetOutsideValue(1)
    io.writeImage(of.GetOutput(), fout)
    return fout


def run(files):
    """run the successive otsu segmentation on the files contained in the
    dictionary files"""
    keys = files.keys()
    keys.sort()
    fcom = 'None'
    gf = {}
    cmds = [ (files[k],utils.getModifiedFileName(files[k],"otsu",".nii.gz")) for k in keys ]
    pool = mp.Pool(mp.cpu_count())
    results = pool.map_async(otsu_seg, cmds)
    data = results.get()
    fcom = data[0]
    for d in data[1:]:
        gf[d] = d
        fcom = combine(fcom, d) 
    
    gf['otsu_combined'] = fcom
    return gf
def main():

    scale = int(sys.argv[1])
    run(scale)     

if __name__ == '__main__':
    main()


